from modelscope import AutoModelForCausalLM, AutoTokenizer
from flask import Flask, request, Response, redirect
import threading


DEFAULT_CKPT_PATH = 'Qwen/Qwen2.5-1.5B-Instruct-GPTQ-Int8'

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_name)


app = Flask(__name__, static_folder="static", static_url_path="")
# 并行度为2
semaphore = threading.Semaphore(3)


@app.route("/chat", methods=["POST"])
def chat():
    with semaphore:
        jso = request.get_json()
        print(jso)
        
        # ai
        messages = jso
        text = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
        )
        model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
        generated_ids = model.generate(
            **model_inputs,
            max_new_tokens=512
        )
        generated_ids = [
            output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
        ]
        response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
        print([response])
        return [response]

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000)
    